from typing import Dict

import numpy as np
from ocr_process.infer.infer_engine import OrtInferSession

from taskinfo.uv_processors import UvPreProcessor, UvPostProcessor


class UvPipeline:
    """UV文档处理管道类"""

    def __init__(self, model_path: str, *, use_cuda: bool = False, use_openvino: bool = False) -> None:
        """初始化UV文档处理管道

        Args:
            model_path: 模型路径
            use_cuda: 是否使用CUDA加速
            use_openvino: 是否使用OpenVINO加速
        """
        config = {
            "use_cuda": use_cuda,
            "use_openvino": use_openvino,
            "model_path": model_path
        }
        self.session = OrtInferSession(config)
        self.preprocessor = UvPreProcessor()
        self.postprocessor = UvPostProcessor()

    def __call__(self, img: np.ndarray) -> Dict[str, np.ndarray]:
        """处理单张图像

        Args:
            img: 输入图像

        Returns:
            处理结果字典
        """
        image, meta = self.preprocessor(img)
        outputs = self.session(image)
        result = self.postprocessor(outputs, meta, img)
        return result
